# -*- coding=utf-8 -*-
from openpyxl import load_workbook
import numpy as np
import copy

# --- define parameter---
data_dir = ''
datas_path = data_dir + 'datas.xlsx'
path_offset = 10


class Node:

    def __init__(self, index, x, y, z, res_type, res_value):
        self.index = int(index)
        self.position = np.array([x, y, z])
        self.res_type = res_type
        self.res_value = res_value
        self.neighbors = []

    def __repr__(self):
        return 'Node {}:{}'.format(self.index, self.position)

    def add_neigh(self, neigh):
        self.neighbors.append(neigh.index)

    def not_neigh(self, neigh):
        index = self.neighbors.count(neigh.index)
        return True if index == 0 else False

    def clear_neigh(self):
        self.neighbors = []

    def generate_neigh(self, total, offset, make_light):
        for e_node in total:
            if self.index != e_node.index & self.not_neigh(e_node):
                res_cost = cal_spend(self, e_node, make_light)
                neighbor_path = self.res_value + offset
                if np.max(res_cost) <= neighbor_path:
                    self.add_neigh(e_node)


class Step:

    def __init__(self, node, parent, food=10, warm=10, spend=0):
        self.node = node.index
        self.position = node.position
        self.spend = spend
        self.parent = parent
        self.index = 0
        self.estimate = 1e30
        self.distance = 0
        self.food = food
        self.warm = warm
        self.cost = self.spend + self.estimate
        self.explored = [0]

    def __repr__(self):
        return 'step {} after {}: start node {} at {}, food left {:.4f}, warm left {:.4f}'.format(
            self.index, self.parent, self.node, self.position, self.food, self.warm)

    def __str__(self):
        return 'node {} at {}, distance {:.2f}, food left {:.4f}, warm left {:.4f}'.format(
            self.node, self.position, self.distance, self.food, self.warm)

    def gen_cost(self):
        self.cost = self.spend + self.estimate

    def set_spend(self, spend):
        self.spend = spend
        self.gen_cost()

    def set_est(self, estimate):
        self.estimate = estimate
        self.gen_cost()

    def get_res(self):
        return np.array([self.warm, self.food])

    def add_res(self, node):
        res = [0, 0]
        res[node.res_type] += node.res_value
        self.warm = self.warm + res[0]
        self.food = self.food + res[1]

    def spend_res(self, parent, spend):
        self.food = parent.food - spend[1]
        self.warm = parent.warm - spend[0]

    def add_explor(self, node_list, node):
        self.explored = copy.deepcopy(node_list)
        self.explored.append(node)


def cal_spend(start, end, light_threshold):
    grad = end.position[2] - start.position[2]
    if grad == 0:
        avg = 5
    else:
        avg = 6 if grad > 0 else 4
    dis = cal_dis(start, end) / 100

    if light_threshold == 0:
        return dis * avg, dis * avg
    else:
        return (dis - 0.2) * avg, (dis * avg) + 0.5


def cal_dis(start, end):
    return np.sqrt(np.sum(np.square(start.position - end.position)))


def contain(node, step_list, total_steps):
    flag = False
    for index in step_list:
        step_i = total_steps[index]
        if node == step_i.node:
            flag = True
    return flag


def get_spend(c_step, esti_type):
    spend_func = None
    if esti_type == 'hop':
        spend_func = hop_spend
    elif esti_type == 'distance':
        spend_func = dis_spend
    elif esti_type == 'make_light':
        spend_func = dis_spend
    return spend_func(c_step)


def hop_spend(c_step):
    return c_step.spend + 1


def dis_spend(c_step):
    return c_step.distance


def get_est(end, final, current_res, esti_type):
    est_func = None
    res_publish = 4

    if esti_type == 'hop':
        est_func = hop_est
    elif esti_type == 'distance':
        est_func = dis_est
    elif esti_type == 'make_light':
        est_func = dis_est
    return est_func(end, final, current_res, res_publish)


def hop_est(end, final, current_res, publish):
    # dis = cal_dis(end, final) / 100
    dis = 0
    min_res = np.where(current_res == np.min(current_res))
    res_value = end.res_value if end.res_type == min_res else end.res_value / publish
    exist_dis = res_value / 4
    return 0 if end.index == final.index else dis - exist_dis


def dis_est(end, final, current_res, publish):
    dis = cal_dis(end, final)
    min_res = np.where(current_res == np.min(current_res))
    res_value = end.res_value if end.res_type == min_res else end.res_value / publish
    exist_dis = res_value / 4
    return 0 if end.index == final.index else dis - exist_dis


def live(food, warm, cost, value):
    return (food - cost[1] > value) & (warm - cost[0] > value)


def lowest_cost(all_nodes):
    cost = 1e30
    select = None
    for index in range(len(all_nodes)):
        node = all_nodes[index]
        if node.cost < cost:
            select = index
            cost = node.cost
    return select


def load_data():
    # --- load datas ---
    datas_sheet = load_workbook(datas_path).active
    nodes = []
    for row in datas_sheet.rows:
        if type(row[0].value) == int:
            para = [col.value for col in row]
            nodes.append(Node(*para))
    return nodes


def search(light_threshold, nodes, est_type):
    # --- search path ---
    steps = [Step(nodes[0], 0)]
    frontier = []
    found = False
    failure = False
    final_node = nodes[-1]

    while True:
        cur_step = steps[-1]
        s_node = nodes[cur_step.node]
        cur_res = np.min(cur_step.get_res())
        s_node.clear_neigh()
        children = s_node.neighbors
        cur_food = cur_step.food
        cur_warm = cur_step.warm
        make_light = (cur_food - cur_warm >= light_threshold) & (est_type == "make_light")
        s_node.generate_neigh(nodes, cur_res, make_light)

        for child_id in children:

            child = nodes[child_id]
            explored = contain(child_id, cur_step.explored, steps)
            if not explored:
                res_spend = cal_spend(s_node, child, make_light)
                dd_value = -5 if child.position[0] != 1000 else -3
                lived = live(cur_food, cur_warm, res_spend, dd_value)
                # far = child.position[0] > s_node.position[0]
                far = True
                if lived & far:
                    next_step = Step(child, cur_step.index)
                    next_step.distance = cal_dis(s_node, child)
                    next_step.set_spend(get_spend(cur_step, est_type))
                    next_step.spend_res(cur_step, res_spend)
                    next_step.set_est(get_est(child, final_node, next_step.get_res(), est_type))
                    next_step.add_res(child)
                    frontier.append(next_step)
                    if child.position[0] == 1000:
                        found = True
                        next_step.index = cur_step.index + 1
                        steps.append(next_step)
                        break

        if found:
            break
        if len(frontier) != 0:
            selected = lowest_cost(frontier)
            current = frontier[selected]
            # print(current)
            frontier.remove(current)
            current.index = cur_step.index + 1
            current.add_explor(steps[current.parent].explored, current.index)
            steps.append(current)
        else:
            failure = True
            break

    right_steps = []
    total_dis = 0
    if not failure:
        final_step = steps[-1]
        while True:
            right_steps.append(final_step)
            total_dis += final_step.distance
            if final_step.index == 0:
                break
            bd_id = final_step.parent
            bd_step = steps[bd_id]
            final_step = bd_step
        right_steps.reverse()

    return ~failure, right_steps, total_dis


estimate_types = ['hop', 'distance', 'make_light']
estimate_type = 2
thresholds = np.arange(0.5, 2, 0.1)
least_step = 1e30
total_results = []
results = None

node_datas = load_data()

for threshold in thresholds:
    success, result_steps, result_distance = search(threshold, node_datas, estimate_types[estimate_type])
    if success:
        total_results.append(result_steps)
        if len(result_steps) < least_step:
            results = (result_steps, result_distance, threshold)

if results is not None:
    if estimate_type == 2:
        print('The right path contain {} steps, with total distance of {:.3f} and light threshold of {}'
              .format(len(results[0]) - 1, results[1], results[2]))
    else:
        print('The right path contain {} steps, with total distance of {:.3f}'
              .format(len(results[0]) - 1, results[1]))
    for step in results[0]:
        print(step)
else:
    print("Cannot found any right path")
